#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
任务完成率批量预测工具
这个脚本可以读取包含多个场景的CSV文件，对每个场景进行预测，并输出结果
"""

import os
import sys
import pandas as pd
import numpy as np
from predict_completion_rate import predict_completion_rate

def batch_predict(input_file, output_file=None):
    """
    批量预测任务完成率
    
    参数:
    input_file (str): 输入CSV文件路径，应包含day_of_week, is_holiday, planned_tasks列
    output_file (str, optional): 输出CSV文件路径，如果不提供则打印到控制台
    
    返回:
    DataFrame: 包含预测结果的数据框
    """
    # 检查输入文件是否存在
    if not os.path.exists(input_file):
        print(f"错误：找不到输入文件 {input_file}")
        return None
    
    # 读取输入文件
    try:
        data = pd.read_csv(input_file)
    except Exception as e:
        print(f"读取输入文件时出错：{e}")
        return None
    
    # 检查必要的列是否存在
    required_columns = ['day_of_week', 'is_holiday', 'planned_tasks']
    missing_columns = [col for col in required_columns if col not in data.columns]
    if missing_columns:
        print(f"错误：输入文件缺少必要的列：{', '.join(missing_columns)}")
        return None
    
    # 创建结果列
    data['predicted_completion_rate'] = np.nan
    
    # 对每行进行预测
    print(f"正在对{len(data)}个场景进行预测...")
    for i, row in data.iterrows():
        predicted_rate = predict_completion_rate(
            row['day_of_week'], 
            row['is_holiday'], 
            row['planned_tasks']
        )
        data.at[i, 'predicted_completion_rate'] = predicted_rate
    
    # 添加建议列
    data['suggestion'] = ''
    for i, row in data.iterrows():
        suggestions = []
        if row['predicted_completion_rate'] < 0.8:
            if row['planned_tasks'] > 9:
                suggestions.append("减少计划任务数量，建议控制在7-9个")
            if row['day_of_week'] in [1, 6, 7]:
                suggestions.append("在周一或周末，需要额外的专注度和动力")
            if row['is_holiday'] == 1:
                suggestions.append("在假日期间，设置更明确的目标和奖励机制")
        data.at[i, 'suggestion'] = '; '.join(suggestions) if suggestions else "无特别建议"
    
    # 保存或打印结果
    if output_file:
        data.to_csv(output_file, index=False)
        print(f"预测结果已保存到 {output_file}")
    else:
        print("\n预测结果:")
        pd.set_option('display.max_rows', None)
        pd.set_option('display.max_columns', None)
        pd.set_option('display.width', 1000)
        print(data)
    
    return data

def create_example_input():
    """创建示例输入文件"""
    example_file = 'data/prediction_scenarios.csv'
    
    # 创建示例数据
    scenarios = []
    
    # 不同星期的场景
    for day in range(1, 8):
        scenarios.append({
            'scenario_name': f'星期{day}，非假日，8个任务',
            'day_of_week': day,
            'is_holiday': 0,
            'planned_tasks': 8
        })
    
    # 假日与非假日对比
    scenarios.append({
        'scenario_name': '周三，假日，8个任务',
        'day_of_week': 3,
        'is_holiday': 1,
        'planned_tasks': 8
    })
    
    # 不同任务数量
    for tasks in [5, 7, 9, 11]:
        scenarios.append({
            'scenario_name': f'周三，非假日，{tasks}个任务',
            'day_of_week': 3,
            'is_holiday': 0,
            'planned_tasks': tasks
        })
    
    # 创建DataFrame并保存
    df = pd.DataFrame(scenarios)
    df.to_csv(example_file, index=False)
    print(f"示例输入文件已创建：{example_file}")
    
    return example_file

def main():
    """主函数"""
    print("\n" + "="*50)
    print("任务完成率批量预测工具")
    print("="*50)
    
    # 检查命令行参数
    if len(sys.argv) > 1:
        input_file = sys.argv[1]
        output_file = sys.argv[2] if len(sys.argv) > 2 else 'output/batch_prediction_results.csv'
    else:
        # 创建并使用示例输入文件
        print("未提供输入文件，将创建示例输入文件...")
        input_file = create_example_input()
        output_file = 'output/batch_prediction_results.csv'
    
    # 确保输出目录存在
    os.makedirs(os.path.dirname(output_file), exist_ok=True)
    
    # 运行批量预测
    batch_predict(input_file, output_file)
    
    print("\n" + "="*50)

if __name__ == "__main__":
    main()